Building a personalized code assistant with open-source LLMs using RAG Fine-tuning

🌈 Abstract

The article discusses the limitations of large language models (LLMs) in code generation and how Retrieval-Augmented Generation (RAG) can address these limitations. It presents a study on fine-tuning the Mistral 7B Instruct v0.2 model using the Together Fine-tuning API and the Morph Code API to improve the model's performance on code generation tasks across different codebases.

🙋 Q&A

[01] Limitations of LLMs in Code Generation

1. What are the two main reasons why LLMs often fall short in code generation?

• Hallucinations: LLMs can generate code that references specific codebases, but the code may not be accurate or up-to-date.
• Outdated knowledge: LLMs may not have the latest coding standards and practices, leading to suboptimal code generation.

2. How does Retrieval-Augmented Generation (RAG) address these limitations? RAG integrates retrieval methods into the text generation process to provide the model with up-to-date and relevant information from external knowledge sources, such as internal documents or code repositories.

[02] RAG Fine-tuning for Code Generation

1. What are the key steps in the RAG fine-tuning process?

• Indexing phase: External knowledge sources (e.g., internal documents) are divided into chunks, transformed into vectors using an embedding model, and stored in a vector database.
• Querying phase: Relevant information is retrieved from the vector database and combined with the initial query in a specific format, which is then used by the generation model to produce the final output.

2. How did the researchers enhance the indexing and retrieval stages of RAG? The researchers partnered with Morph Labs to leverage their advanced technologies in codebase search and synthetic data generation.

3. How did the researchers fine-tune the generation model (Mistral 7B Instruct v0.2) to improve its performance? The researchers fine-tuned the Mistral 7B Instruct v0.2 model using Together's Fine-tuning API to ensure the model's responses are up-to-date by incorporating the latest coding standards and practices.

[03] Experimental Setup and Results

1. What codebases were used in the experiments, and what was the purpose of selecting these codebases? The experiments were conducted on five different codebases: Axolotl, Deepspeed, vLLM, Mapbox, and WandB. These codebases offer comprehensive coverage of commonly used machine learning codebases spanning training, inference, and interface. The researchers chose these codebases because many existing LLMs show limited understanding of them due to their complexity and rapidly evolving nature.

2. What evaluation metric was used to assess the performance of the generated examples? The researchers used the HitRate (%) as the evaluation metric. For each query, they checked whether a key identifier string (labeled by human annotators) was present in the generated code, indicating the model's understanding and reasoning capability for the codebase.

3. How did the performance of the RAG fine-tuned Mistral 7B Instruct v0.2 model compare to other models (Mistral 7B Instruct v0.2, Claude 3 Opus, and GPT-4o)? The RAG fine-tuned Mistral 7B Instruct v0.2 model achieved significant performance improvements over the original model on 4 out of 5 codebases under the RAG setting. The RAG fine-tuned models with 7 billion parameters even achieved at par or better performance compared to much larger models like GPT-4o and Claude 3 Opus, except on the WandB codebase.

4. What additional benefits did the researchers find in using the RAG fine-tuned models? Beyond the quality improvements, the RAG fine-tuned models offer significant economic advantages, being up to 150x cheaper to build and deploy on the Together platform, and 3.7x faster during inference compared to the larger models.